# !/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time        : 2020/12/5 20:25
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : DirectiveSolveLinearSystemEquations.py
@Version     : Version 1.0.0
@Description : TODO 自己实现(选主元和不选主元)高斯消元法,直接三角(选主元和不选主元)分解法，矩阵LU分解，矩阵PLU分解,两种平方根法,追赶法
@Created By  : PyCharm
"""
import numpy as np
from scipy.linalg import det, lu, lu_solve, lu_factor, inv, eigvals


# 自己原创
def lu_decomposition(coefficient_matrix: np.ndarray, right_hand_side_vector: np.ndarray):
    """
    实现方程Ax=b系数矩阵A的LU decomposition
    :param coefficient_matrix: 初始系数矩阵A
    :param right_hand_side_vector: 初始常数列向量b
    :return: 单位下三角矩阵L，上三角矩阵U，常数列向量b
    """
    # first step: evaluate the lu decomposition condition
    rows, columns = coefficient_matrix.shape
    if rows == columns:  # judge if it is a square matrix
        for k in range(rows):  # 判断各阶顺序主子式是否为0
            if det(coefficient_matrix[:k + 1, :k + 1]) == 0:
                raise Exception("cannot generate LU decomposition")
        else:
            # LU decomposition
            lower_triangular_matrix = np.eye(rows)
            for k in range(rows - 1):
                for row in range(k + 1, rows):
                    multiplier = coefficient_matrix[row, k] / coefficient_matrix[k, k]
                    coefficient_matrix[row, k:] += -multiplier * (coefficient_matrix[k, k:])
                    right_hand_side_vector[row] += -multiplier * right_hand_side_vector[k]
                    lower_triangular_matrix[row, k] = multiplier
    else:
        raise Exception("ERROR:please pass a square matrix.")

    return lower_triangular_matrix, coefficient_matrix, right_hand_side_vector


# 自己原创
def pivot_lu_decomposition(coefficient_matrix: np.ndarray, right_hand_side_vector: np.ndarray):
    """
    实现方程Ax=b系数矩阵A的pivoted LU decomposition
    :param coefficient_matrix: 初始系数矩阵A
    :param right_hand_side_vector: 初始常数列向量b
    :return: 排列矩阵P，单位下三角矩阵L，上三角矩阵U，常数列向量b
    """
    # first step: evaluate the lu decomposition condition
    rows, columns = coefficient_matrix.shape
    if rows == columns:  # judge if it is a square matrix
        for k in range(rows):  # 判断各阶顺序主子式是否为0
            if det(coefficient_matrix[:k + 1, :k + 1]) == 0:
                raise Exception("cannot generate LU decomposition")
        else:
            # pivot LU decomposition
            lower_triangular_matrix = np.eye(rows)
            permutation_matrix = np.eye(rows)
            for k in range(rows - 1):
                max_index = np.argmax(abs(coefficient_matrix[k:, k]))
                coefficient_matrix[[k, max_index + k], :] = coefficient_matrix[[max_index + k, k], :]
                right_hand_side_vector[[k, max_index + k], :] = right_hand_side_vector[[max_index + k, k], :]
                permutation_matrix[[k, max_index + k], :] = permutation_matrix[[max_index + k, k], :]
                for row in range(k + 1, rows):
                    multiplier = coefficient_matrix[row, k] / coefficient_matrix[k, k]
                    coefficient_matrix[row, k:] += -multiplier * (coefficient_matrix[k, k:])
                    right_hand_side_vector[row] += -multiplier * right_hand_side_vector[k]
                    lower_triangular_matrix[row, k] = multiplier
    else:
        raise Exception("ERROR:please pass a square matrix.")

    return permutation_matrix, lower_triangular_matrix, coefficient_matrix, right_hand_side_vector


# 自己原创
def gaussian_elimination(coefficient_matrix: np.ndarray,
                         right_hand_side_vector: np.ndarray,
                         decomposition_function=pivot_lu_decomposition):
    """
    实现高斯消元法解方程Ax=b
    :param decomposition_function:分解函数
    :param coefficient_matrix: 初始系数矩阵A
    :param right_hand_side_vector: 初始列向量b
    :return: 解向量x
    """
    *_, upper_triangular_matrix, constant_column_vector = decomposition_function(coefficient_matrix,
                                                                                 right_hand_side_vector)
    rows = right_hand_side_vector.shape[0]
    if upper_triangular_matrix.shape[0] == rows:  # 判断上三角矩阵的行数是否等于列向量行数
        # 此处不指定双精度浮点数据类型，则可能会损失精度，导致结果偏差太大
        variable_column_vector = np.ones_like(right_hand_side_vector, dtype=np.float64)  # 构建一个与列向量同形状的1矩阵
        for row in range(rows - 1, -1, -1):
            prod = 0  # 声明累积变量，初始化为0
            for column in range(row + 1, rows):
                prod += variable_column_vector[column] * upper_triangular_matrix[row, column]
            # 回代运算得到解向量
            variable_column_vector[row] = (constant_column_vector[row] - prod) / upper_triangular_matrix[row, row]
        return variable_column_vector
    else:
        raise Exception("系数矩阵的行数与常数列向量行数不相等")


# 自己原创
def triangle_decomposition(coefficient_matrix: np.ndarray):
    """
    实现方程Ax=b系数矩阵A的Doolittle decomposition
    :param coefficient_matrix: 初始系数矩阵A
    :return: 单位下三角矩阵L,上三角矩阵U
    """
    # first step: evaluate the decomposition condition
    rows, columns = coefficient_matrix.shape
    if rows == columns:  # judge if it is a square matrix
        for k in range(rows):  # 判断各阶顺序主子式是否为0
            if det(coefficient_matrix[:k + 1, :k + 1]) == 0:
                raise Exception("cannot decompose ")
        else:
            # directive triangle  decomposition
            # 初始化
            square_ones_matrix = np.ones((rows, columns))
            lower_triangle_matrix = np.tril(square_ones_matrix)
            upper_triangle_matrix = np.triu(square_ones_matrix)
            # 计算第1行的U(1,i),第1列的L(i,1)
            upper_triangle_matrix[0, :] = coefficient_matrix[0, :]
            lower_triangle_matrix[1:, 0] = coefficient_matrix[1:, 0] / upper_triangle_matrix[0, 0]
            # 计算第row行的U(row,i),第row列的L(i,row)
            for row in range(1, rows - 1):

                for i in range(row, rows):
                    # 定义一个累积和变量
                    row_i_prod = 0
                    for k in range(row):
                        row_i_prod += lower_triangle_matrix[row, k] * upper_triangle_matrix[k, i]
                    upper_triangle_matrix[row, i] = coefficient_matrix[row, i] - row_i_prod

                for i in range(row + 1, rows):
                    # 定义一个累积和变量
                    i_row_prod = 0
                    for k in range(row):
                        i_row_prod += lower_triangle_matrix[i, k] * upper_triangle_matrix[k, row]
                    lower_triangle_matrix[i, row] = (coefficient_matrix[i, row] - i_row_prod) / upper_triangle_matrix[
                        row, row]

            # 定义一个累积和变量,单独计算n行n列的U(n,n)
            row_i_prod = 0
            for k in range(rows - 1):
                row_i_prod += lower_triangle_matrix[rows - 1, k] * upper_triangle_matrix[k, rows - 1]
            upper_triangle_matrix[rows - 1, rows - 1] = coefficient_matrix[rows - 1, rows - 1] - row_i_prod

            return lower_triangle_matrix, upper_triangle_matrix
    else:
        raise Exception("ERROR:please pass a square matrix.")


# 暂时没有实现
def pivot_triangle_decomposition():
    pass


# 自己原创
def directive_triangle_decomposition(coefficient_matrix: np.ndarray, right_hand_side_vector: np.ndarray,
                                     decomposition_function=triangle_decomposition):
    """
    实现方程Ax=b系数矩阵A directive triangle decomposition
    :param decomposition_function: 三角分解函数
    :param coefficient_matrix: 初始系数矩阵A
    :param right_hand_side_vector: 初始常数列向量b
    :return: 解向量x
    """
    l, u = decomposition_function(coefficient_matrix)
    rows, columns = coefficient_matrix.shape
    # Ly=b，此处不指定双精度浮点数据类型，则可能会损失精度，导致结果偏差太大
    y = np.ones_like(right_hand_side_vector, dtype=np.float64)
    y[0, 0] = right_hand_side_vector[0, 0]
    for row in range(1, rows):
        # 定义一个累积和变量
        i_k_prod = 0
        for k in range(row):
            i_k_prod += l[row, k] * y[k, 0]
        y[row, 0] = right_hand_side_vector[row, 0] - i_k_prod

    # Ux=y，此处不指定双精度浮点数据类型，则可能会损失精度，导致结果偏差太大
    x = np.ones_like(right_hand_side_vector, dtype=np.float64)
    x[rows - 1] = y[rows - 1] / u[rows - 1, rows - 1]
    for row in range(rows - 2, 0, -1):
        i_k_prod = 0
        for k in range(row + 1, rows):
            i_k_prod += u[row, k] * x[k, 0]
        x[row, 0] = (y[row, 0] - i_k_prod) / u[row, row]

    return x


# 自己原创平方根法
def cholesky_decomposition(sym_pos_define_matrix: np.ndarray, right_hand_side_vector: np.ndarray):
    rows, columns = sym_pos_define_matrix.shape
    # judge if it is a square matrix
    if rows == columns:
        # 判断是否对称
        if (sym_pos_define_matrix.T == sym_pos_define_matrix).all():
            for k in range(rows):  # 判断各阶顺序主子式是否为0,即判定是否正定
                if det(sym_pos_define_matrix[:k + 1, :k + 1]) <= 0:
                    raise Exception("cannot decompose")
            else:
                # 初始化单位下三角矩阵L
                square_ones_matrix = np.ones((rows, columns))
                lower_triangle_matrix = np.tril(square_ones_matrix)
                for j in range(rows):
                    prod = 0
                    for k in range(j):
                        prod += lower_triangle_matrix[j, k] * lower_triangle_matrix[j, k]
                    lower_triangle_matrix[j, j] = np.sqrt(sym_pos_define_matrix[j, j] - prod)

                    for i in range(j + 1, rows):
                        prod = 0
                        for k in range(j):
                            prod += lower_triangle_matrix[i, k] * lower_triangle_matrix[j, k]
                        lower_triangle_matrix[i, j] = (sym_pos_define_matrix[i, j] - prod) / lower_triangle_matrix[j, j]

        else:
            raise Exception("error,the input matrix must be a symmetric-matrix")
    else:
        raise Exception("ERROR:please pass a square matrix.")
    return inv(lower_triangle_matrix.transpose()) @ inv(lower_triangle_matrix) @ right_hand_side_vector


# 自己原创改进平方根
def improved_square_root(symmetric_positive_define_matrix: np.ndarray, right_hand_side_vector: np.ndarray):
    rows, columns = symmetric_positive_define_matrix.shape
    # judge if it is a square matrix
    if rows == columns:
        # 判断是否对称
        if (symmetric_positive_define_matrix.T == symmetric_positive_define_matrix).all():
            for k in range(rows):  # 判断各阶顺序主子式是否为0,即判定是否正定
                if det(symmetric_positive_define_matrix[:k + 1, :k + 1]) == 0:
                    raise Exception("error, cannot decompose")
            else:
                # 初始化单位下三角矩阵L
                square_ones_matrix = np.ones((rows, columns))
                lower_triangle_matrix = np.tril(square_ones_matrix)
                diagonal_matrix = np.diag(np.diag(square_ones_matrix))
                t = lower_triangle_matrix @ diagonal_matrix
                diagonal_matrix[0, 0] = symmetric_positive_define_matrix[0, 0]
                for i in range(1, rows):
                    for j in range(i):
                        prod = 0
                        for k in range(j):
                            prod += t[i, k] * lower_triangle_matrix[j, k]
                        t[i, j] = symmetric_positive_define_matrix[i, j] - prod
                        lower_triangle_matrix[i, j] = t[i, j] / diagonal_matrix[j, j]
                    prod = 0
                    for k in range(i):
                        prod += t[i, k] * lower_triangle_matrix[i, k]
                    diagonal_matrix[i, i] = symmetric_positive_define_matrix[i, i] - prod

                return inv(
                    lower_triangle_matrix @ diagonal_matrix @ lower_triangle_matrix.T) @ right_hand_side_vector
                # 自己写的发现第一个解x0=1.,而不是1.11111111，有精度损失，所以采用矩阵运算
                # # Ly=b,此处不指定双精度浮点数据类型，则可能会损失精度，导致结果偏差太大
                # y = np.ones_like(right_hand_side_vector, dtype=np.float64)
                # for i in range(rows):
                #     prod = 0
                #     for k in range(i):
                #         prod += lower_triangle_matrix[i, k] * y[k, 0]
                #     y[i, 0] = right_hand_side_vector[i, 0] - prod
                # # DL^Tx=y，此处不指定双精度浮点数据类型，则可能会损失精度，导致结果偏差太大
                # x = np.ones_like(right_hand_side_vector, dtype=np.float64)
                # x[rows - 1, 0] = y[rows - 1, 0] / diagonal_matrix[rows - 1, rows - 1]
                # for i in range(rows - 2, 0, -1):
                #     prod = 0
                #     for k in range(i + 1, rows):
                #         prod += lower_triangle_matrix[k, i] * x[k, 0]
                #     x[i, 0] = y[i, 0] / diagonal_matrix[i, i] - prod
                # return x

                # for i in range(rows):
                #     for j in range(i):
                #         prod = 0
                #         for k in range(j):
                #             prod += lower_triangle_matrix[i, k] * diagonal_matrix[k, k] * lower_triangle_matrix[j, k]
                #         lower_triangle_matrix[i, j] = (symmetric_positive_define_matrix[i, j] - prod) / diagonal_matrix[
                #             j, j]
                #     prod = 0
                #     for k in range(i):
                #         prod += lower_triangle_matrix[i, k] * diagonal_matrix[k, k] * lower_triangle_matrix[i, k]
                #     diagonal_matrix[i, i] = symmetric_positive_define_matrix[i, i] - prod
        else:
            raise Exception("error,the input matrix must be a symmetric-matrix")
    else:
        raise Exception("ERROR:please pass a square matrix.")


# 自己原创，注意这里没有做矩阵类型判断，该方法仅适用于解系数矩阵为对角占优的三对角线方程组,P159
def chasing_thomas(triple_diagonal_matrix: np.ndarray, f_vector: np.ndarray):
    b_diagonal = np.diag(triple_diagonal_matrix)
    c_diagonal = np.diag(triple_diagonal_matrix, k=1)
    a_diagonal = np.diag(triple_diagonal_matrix, k=-1)
    # 计算Bata_i,dtype=np.float64,提高运算精度，否则默认最小精度
    beta_diagonal = np.ones_like(c_diagonal, dtype=np.float64)
    beta_diagonal[0] = c_diagonal[0] / b_diagonal[0]
    rows = triple_diagonal_matrix.shape[0]
    for i in range(1, rows - 1):
        beta_diagonal[i] = c_diagonal[i] / (b_diagonal[i] - a_diagonal[i] * beta_diagonal[i - 1])
    # 解Ly=f
    y = np.ones_like(f_vector, dtype=np.float64)
    y[0, 0] = f_vector[0, 0] / b_diagonal[0]
    for i in range(1, rows):
        y[i, 0] = (f_vector[i, 0] - a_diagonal[i - 1] * y[i - 1, 0]) / (
                b_diagonal[i] - a_diagonal[i - 1] * beta_diagonal[i - 1])
    # 解Ux=y,自己写的发现第一个解x0=1.,而不是0.8333333333333334，有精度损失，
    x = np.ones_like(f_vector, dtype=np.float64)
    x[rows - 1, 0] = y[rows - 1, 0]
    for i in range(rows - 2, 0, -1):
        x[i, 0] = y[i, 0] - beta_diagonal[i] * x[i + 1, 0]
    return x


def vector_norm(vector: np.ndarray, p=None):
    """
    计算向量的p-范数
    :param vector: 实向量或者复向量
    :param p: 指定类型的范数，默认是oo范数
    :return: 指定向量范数和p值
    """
    if p is None:
        return abs(vector).max(), p
    elif p >= 1:
        return np.power(np.sum(np.power(abs(vector), p)), 1 / p), p
    else:
        raise Exception("error,p must be an integer , greater than  or equal to 1")


def matrix_norm(square_matrix: np.ndarray, v=None):
    """
    计算矩阵范数，其中v=1,2,oo
    :param square_matrix: 实数方阵或者复数方阵
    :param v: 范数类型，默认是弗罗贝尼乌斯范数
    :return: 指定类型范数
    """
    if square_matrix.shape[0] == square_matrix.shape[1]:
        if v is None:
            # 弗罗贝尼乌斯范数
            return np.power(np.sum(np.power(square_matrix, 2)), 1 / 2)
        elif v == 0:
            # 矩阵的列范数
            return np.max(np.sum(abs(square_matrix), v))
        elif v == 1:
            # 矩阵的行范数
            return np.max(np.sum(abs(square_matrix), v))
        elif v == 2:
            # 矩阵的2-范数
            return np.sqrt(eigvals(square_matrix.T @ square_matrix).max()).real
        else:
            raise Exception("sorry,currently,v must be None,0,1 or2")
    else:
        raise Exception("error,input matrix must be a square matrix.")


def condition_number(square_matrix: np.ndarray, v=None):
    """
    计算矩阵条件数，其中v=oo,1,2
    :param square_matrix:实数方阵或者复数方阵
    :param v:范数类型，默认是oo范数
    :return:指定范数下的矩阵条件数
    """
    if square_matrix.shape[0] == square_matrix.shape[1]:
        if det(square_matrix) != 0:
            if v is None:
                return matrix_norm(square_matrix, v=1) * matrix_norm(inv(square_matrix), v=1)
            elif v == 1:
                return matrix_norm(square_matrix, v=0) * matrix_norm(inv(square_matrix), v=0)
            elif v == 2:
                eigenvalue_fraction = matrix_norm(square_matrix, v=2) / np.sqrt(
                    eigvals(square_matrix.T @ square_matrix).min()).real
                return np.sqrt(eigenvalue_fraction)
            else:
                raise Exception("sorry,currently,v must be None,1 or2")
        else:
            raise Exception("error,input matrix must be a non-singular matrix")
    else:
        raise Exception("error,input matrix must be a square matrix.")


if __name__ == '__main__':
    """
    # 改进平方根法测试成功，来源详见李庆扬数值分析第5版P177，e.g.10
    symmetric_positive_define_matrix1 = np.array([[2, -1, 1], [-1, -2, 3], [1, 3, 1]])
    column_vector = np.array([4, 5, 6]).reshape((3, 1))
    print(improved_square_root(symmetric_positive_define_matrix1, column_vector))
    """
    """
    # 追赶法测试成功，来源详见李庆扬数值分析第5版P177，e.g.9
    # dtype指定为双精度浮点提高运算精度
    tri_diagonal = np.array([[2, -1, 0, 0, 0],
                             [-1, 2, -1, 0, 0],
                             [0, -1, 2, -1, 0],
                             [0, 0, -1, 2, -1],
                             [0, 0, 0, -1, 2]], dtype=np.float64)
    f = np.array([1, 0, 0, 0, 0], dtype=np.float64).reshape((5, 1))
    print(chasing_thomas(tri_diagonal, f))
    """
    # 楚列斯基分解测试1，通过
    # symmetric_positive_define_matrix0 = np.array([[1, 1, 1], [1,2, 3], [1, 3, 6]])
    # 楚列斯基分解测试2，通过
    # symmetric_positive_define_matrix1 = np.array([[1, 1, 1], [1,4,2], [1, 2, 8]])
    """
    symmetric_positive_define_matrix2 = np.array([[4, 2, -4, 0, 2, 4, 0, 0],
                                                  [2, 2, -1, -2, 1, 3, 2, 0],
                                                  [-4, -1, 14, 1, -8, -3, 5, 6],
                                                  [0, -2, 1, 6, -1, -4, -3, 3],
                                                  [2, 1, -8, -1, 22, 4, -10, -3],
                                                  [4, 3, -3, -4, 4, 11, 1, -4],
                                                  [0, 2, 5, -3, -10, 1, 14, 2],
                                                  [0, 0, 6, 3, -3, -4, 2, 19]],dtype=np.float64)
    column_vector = np.array([0, -6, 20, 23, 9, -22, -15, 45],dtype=np.float64).reshape((8, 1))
    print(cholesky_decomposition(symmetric_positive_define_matrix2, column_vector))
    """
    # 直接三角分解法-不选主元三角分解法测试成功，来源详见李庆扬数值分析第5版P153，e.g.5
    """ 
    square_matrix = np.array([[1, 2, 3], [2, 5, 2], [3, 1, 5]])
    column_vector = np.array([14, 18, 20]).reshape((3, 1))
    print(directive_triangle_decomposition(square_matrix, column_vector))
    """
    # 第一组测试方程，来源详见李庆扬数值分析第5版P148-149
    """
    square_matrix = np.array([[0.001, 2.000, 3.000], [-1.000, 3.712, 4.623], [-2.000, 1.072, 5.643]],
    dtype=np.float64)
    vector = np.array([1.000, 2.000, 3.000], dtype=np.float64).reshape((3, 1))
    """
    # 第二组测试方程，来源详见李庆扬数值分析第5版P143，148
    """square_matrix = np.array([[1, 1, 1], [0, 4, -1], [2, -2, 1]], dtype=np.float64)
    print("原系数矩阵为:\n{}".format(square_matrix))
    vector = np.array([6, 5, 1], dtype=np.float64).reshape(3, 1)
    print("原常数列向量为:\n{}".format(vector))
    """
    # 为避免相互影响，这里单独复制系数矩阵和常数列向量，分别测试系统函数和自己写的函数
    # 第三组测试方程，来源详见李庆扬数值分析第5版P178 e.g.1
    """
    square_matrix = np.array([[10, -7, 0, 1], [-3, 2.099999, 6, 2], [5, -1, 5, -1], [2, 1, 0, 2]])
    vector = np.array([8, 5.900001, 5, 1]).reshape((4, 1))
    copy_square_matrix1 = square_matrix.copy()
    copy_vector1 = vector.copy()
    copy_square_matrix2 = square_matrix.copy()
    copy_vector2 = vector.copy()
    solution = lu_solve(lu_factor(square_matrix), vector)
    print("系统函数LU分解实现的解向量x=\n{}".format(solution))
    l,u,_ = lu_decomposition(copy_square_matrix1, copy_vector1)
    print("Ax=b中系数矩阵A=LU分解的矩阵\nL=\n{}\nU=\n{}".format(l, u))
    solution_vector = gaussian_elimination(copy_square_matrix2, copy_vector2)
    print("自己写的不选主元高斯消去法解向量x=\n{}".format(solution_vector))
    print("detA={}".format(det(copy_square_matrix2)))
    """
    # 第四组测试方程，来源详见李庆扬数值分析第5版P178 e.g.2

    #     square_mat = np.array([[3.01, 6.03, 1.99], [1.27, 4.16, -1.23], [0.987, -4.81, 9.34]])
    #     vector = np.array([1, 1, 1], dtype=np.float64).reshape((3, 1))
    #     print("A=\n{}\nb=\n{}".format(square_mat, vector))
    #     print("detA={}".format(det(square_mat)))
    #     copy_square_matrix = square_mat.copy()
    #     copy_vector = vector.copy()
    #     solution = lu_solve(lu_factor(square_mat), vector)
    #     print("系统函数LU分解实现的解向量x=\n{}".format(solution))
    #     solution_vector = gaussian_elimination(copy_square_matrix, copy_vector)
    #     print("自己写的列主元高斯消去法解向量x=\n{}".format(solution_vector))
    #     print("""矩阵A
    # oo条件数cond(A)oo=
    # {}
    # 谱条件数cond(A)2=
    # {}""".format(condition_number(copy_square_matrix), condition_number(copy_square_matrix, v=2)))

    """
    # 自己写的列主元LU分解
    p1, l1, u1, _ = pivot_lu_decomposition(square_matrix, vector)
    # print(copy_square_matrix, p @ l @ u, sep='\n')
    print(p1, l1, u1, sep='\n')
    # 系统自带LU分解函数
    p2, l2, u2 = lu(copy_square_matrix)
    print(p2, l2, u2, sep='\n')
    """

    """来源详见李庆扬数值分析第5版P145,测试无法LU分解的某些简单矩阵情况
    a = np.array([[0, 1], [1, 0]])
    lu_decomposition(a)
    """
